#    Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.


from collections import OrderedDict
import shutil
from batchgenerators.utilities.file_and_folder_operations import *


if __name__ == "__main__":
    base = "/media/chyang/data/dataset/project/nnformer/nnFormer_raw/nnFormer_raw_data/"

    task_id = 17
    task_name = "Task02_Synapse"
    prefix = 'ABD'
    train_frac = 0.8  # 测试集比例


    out_base = join(base, task_name)
    imagestr = join(out_base, "imagesTr")
    imagests = join(out_base, "imagesTs")
    labelstr = join(out_base, "labelsTr")

    train_patient_names = []
    test_patient_names = []
    train_patient_names = subfiles(imagestr, join=False, suffix = 'nii.gz')
    test_patient_names = subfiles(imagests, join=False, suffix=".nii.gz")
 

    json_dict = OrderedDict()
    json_dict['name'] = "AbdominalOrganSegmentation"
    json_dict['description'] = "Multi-Atlas Labeling Beyond the Cranial Vault Abdominal Organ Segmentation"
    json_dict['tensorImageSize'] = "3D"
    json_dict['reference'] = "https://www.synapse.org/#!Synapse:syn3193805/wiki/217789"
    json_dict['licence'] = "see challenge website"
    json_dict['release'] = "0.0"
    json_dict['modality'] = {
        "0": "CT",
    }
    json_dict['labels'] = OrderedDict({
        "00": "background",
        "01": "spleen",
        "02": "right kidney",
        "03": "left kidney",
        "04": "gallbladder",
        "05": "esophagus",
        "06": "liver"}
    )
    json_dict['numTraining'] = len(train_patient_names)
    json_dict['numTest'] = len(test_patient_names)
    trainset_numer = int(len(train_patient_names) * train_frac) # 测试集图像个数
    json_dict['training'] = [{'image': f"{imagestr}/imagesTr/{train_patient_name}", "label": f"{labelstr}/labelsTr/{train_patient_name[:-12]}.nii.gz"} for i, train_patient_name in enumerate(train_patient_names[:trainset_numer])]
    json_dict['validation'] = [{'image': f"{imagestr}/imagesTr/{train_patient_name}", "label": f"{labelstr}/labelsTr/{train_patient_name[:-12]}.nii.gz"} for i, train_patient_name in enumerate(train_patient_names[trainset_numer:])]
    json_dict['test'] = ["{imagests}/imagesTs/%s" % test_patient_name for test_patient_name in test_patient_names]

    save_json(json_dict, os.path.join(out_base, "dataset.json"))
